"""
数据源适配器模块
支持Excel、JSON、MySQL等多种数据源
实现统一的数据读取接口
"""

from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional
import pandas as pd
import json
import pymysql
from pathlib import Path


class DataSource(ABC):
    """数据源抽象基类"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.validate_config()
    
    @abstractmethod
    def validate_config(self):
        """验证配置有效性"""
        pass
    
    @abstractmethod
    def read_data(self) -> List[Dict[str, Any]]:
        """
        读取数据
        
        Returns:
            数据列表，每条数据为字典格式
        """
        pass


class ExcelDataSource(DataSource):
    """Excel数据源"""
    
    def validate_config(self):
        """验证Excel配置"""
        if 'file_path' not in self.config:
            raise ValueError("Excel数据源配置缺少file_path字段")
        
        file_path = Path(self.config['file_path'])
        if not file_path.exists():
            raise FileNotFoundError(f"Excel文件不存在: {file_path}")
        
        if file_path.suffix not in ['.xlsx', '.xls']:
            raise ValueError(f"不支持的Excel格式: {file_path.suffix}")
    
    def read_data(self) -> List[Dict[str, Any]]:
        """读取Excel数据"""
        file_path = self.config['file_path']
        sheet_name = self.config.get('sheet_name', 0)  # 默认第一个sheet
        
        try:
            # 读取Excel文件
            df = pd.read_excel(file_path, sheet_name=sheet_name)
            
            # 处理空值，替换为None
            df = df.where(pd.notna(df), None)
            
            # 转换为字典列表
            data = df.to_dict('records')
            
            return data
            
        except Exception as e:
            raise ValueError(f"读取Excel文件失败: {str(e)}")


class JsonDataSource(DataSource):
    """JSON数据源"""
    
    def validate_config(self):
        """验证JSON配置"""
        if 'file_path' not in self.config:
            raise ValueError("JSON数据源配置缺少file_path字段")
        
        file_path = Path(self.config['file_path'])
        if not file_path.exists():
            raise FileNotFoundError(f"JSON文件不存在: {file_path}")
        
        if file_path.suffix != '.json':
            raise ValueError(f"不支持的JSON格式: {file_path.suffix}")
    
    def read_data(self) -> List[Dict[str, Any]]:
        """读取JSON数据"""
        file_path = self.config['file_path']
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 如果是单个对象，转换为列表
            if isinstance(data, dict):
                return [data]
            elif isinstance(data, list):
                return data
            else:
                raise ValueError("JSON数据必须是对象或数组格式")
                
        except json.JSONDecodeError as e:
            raise ValueError(f"JSON格式错误: {str(e)}")
        except Exception as e:
            raise ValueError(f"读取JSON文件失败: {str(e)}")


class MysqlDataSource(DataSource):
    """MySQL数据源"""
    
    def validate_config(self):
        """验证MySQL配置"""
        required_fields = ['host', 'port', 'user', 'password', 'database', 'sql']
        for field in required_fields:
            if field not in self.config:
                raise ValueError(f"MySQL数据源配置缺少{field}字段")
    
    def read_data(self) -> List[Dict[str, Any]]:
        """读取MySQL数据"""
        connection = None
        try:
            # 建立数据库连接
            connection = pymysql.connect(
                host=self.config['host'],
                port=int(self.config.get('port', 3306)),
                user=self.config['user'],
                password=self.config['password'],
                database=self.config['database'],
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor  # 使用字典游标
            )
            
            with connection.cursor() as cursor:
                # 执行SQL查询
                sql = self.config['sql']
                cursor.execute(sql)
                
                # 获取所有结果
                data = cursor.fetchall()
                
                return data
                
        except pymysql.Error as e:
            raise ValueError(f"MySQL查询失败: {str(e)}")
        finally:
            if connection:
                connection.close()
    
    def test_connection(self) -> bool:
        """测试数据库连接"""
        try:
            connection = pymysql.connect(
                host=self.config['host'],
                port=int(self.config.get('port', 3306)),
                user=self.config['user'],
                password=self.config['password'],
                database=self.config['database'],
                charset='utf8mb4'
            )
            connection.close()
            return True
        except pymysql.Error:
            return False


class DataSourceFactory:
    """数据源工厂类"""
    
    # 数据源类型映射
    _source_types = {
        'excel': ExcelDataSource,
        'json': JsonDataSource,
        'mysql': MysqlDataSource,
    }
    
    @classmethod
    def create(cls, source_type: str, config: Dict[str, Any]) -> DataSource:
        """
        创建数据源实例
        
        Args:
            source_type: 数据源类型（excel/json/mysql）
            config: 数据源配置
            
        Returns:
            数据源实例
            
        Raises:
            ValueError: 不支持的数据源类型
        """
        source_type = source_type.lower()
        
        if source_type not in cls._source_types:
            raise ValueError(f"不支持的数据源类型: {source_type}")
        
        source_class = cls._source_types[source_type]
        return source_class(config)
    
    @classmethod
    def get_supported_types(cls) -> List[str]:
        """获取支持的数据源类型列表"""
        return list(cls._source_types.keys())
